Chapter 10.3 - Autoregressive Generation Loop

---

📝 Chapter 10: Text Generation

Now that we have a trained (or pretrained) model, how does it actually "speak"?

Generating text is an iterative process. A language model doesn't output an entire sentence at once; it predicts the next single token. We then take that token, add it to our input, and ask the model to predict the next one.

This process is handled by f_token_generator.py and g_text_generator.py.

---

🎲 1. The Core Loop: Generating Tokens

In f_token_generator.py, the generate_next_tokens() function is the heart of inference.

for _ in range(max_new_tokens):
    # 1. Get logits for the current context
    logits = model(idx_cond)[:, -1, :] 

    # 2. Pick the next token
    idx_next = torch.argmax(logits, dim=-1, keepdim=True)

    # 3. Append to the sequence and repeat
    idx = torch.cat((idx, idx_next), dim=1)

Advanced Sampling Methods

Always picking the highest probability token (torch.argmax) is called greedy decoding. It works, but it can sound repetitive or boring. To make the model more creative, we add two techniques:

1. Top-K Sampling: Instead of considering the entire vocabulary, we only look at the top $K$ most likely tokens. The rest of the tokens get their probabilities forced to negative infinity (-inf), meaning they have zero chance of being picked.

2. Temperature Scaling: We divide the logits by a temperature value.

   • temperature < 1.0 makes the peaks sharper (more confident, less random).
   • temperature > 1.0 flattens the distribution (more random, more creative).

After applying temperature, we convert logits to probabilities using Softmax, and then use torch.multinomial() to randomly sample the next token based on those probabilities.

---

🗣️ 2. The User Interface: Text Generator

In g_text_generator.py, we wrap the token generation logic into a user-friendly pipeline.

The generate_text() function does three things:

1. Encode: Uses the tokenizer to turn your string prompt into token IDs.
2. Generate: Calls generate_next_tokens() to get a long list of output token IDs.
3. Decode: Converts those output token IDs back into human-readable text.

Eval Mode

Before generating text, we must set model.eval(). This turns off training-specific layers like Dropout, ensuring our generation is deterministic and uses the full strength of the network. We also wrap the generation in torch.no_grad() to save memory, since we aren't calculating gradients for backpropagation.

---

💻 Code Implementation

Here is the exact PyTorch implementation for the concepts discussed above:

# this takes in the input tensor like tensor([[6109, 3626, 6100, 345]]) and generates the new token 
# and emits  the complte new tensor like tensor([[6109, 3626, 6100, 345, 257]])
import torch
import torch.nn as nn
def generate_next_tokens(model, idx, max_new_tokens, context_size, temperature=0.0, top_k=None, eos_id=None):

    # For-loop is the same as before: Get logits, and only focus on last time step
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -context_size:]
        with torch.no_grad():
            logits = model(idx_cond)
        logits = logits[:, -1, :]

        # New: Filter logits with top_k sampling
        if top_k is not None:
            # Keep only top_k values
            top_logits, _ = torch.topk(logits, top_k)
            min_val = top_logits[:, -1]
            logits = torch.where(logits < min_val, torch.tensor(float("-inf")).to(logits.device), logits)

        # New: Apply temperature scaling
        if temperature > 0.0:
            logits = logits / temperature

            # New (not in book): numerical stability tip to get equivalent results on mps device
            # subtract rowwise max before softmax
            logits = logits - logits.max(dim=-1, keepdim=True).values
            
            # Apply softmax to get probabilities
            probs = torch.softmax(logits, dim=-1)  # (batch_size, context_len)

            # Sample from the distribution
            idx_next = torch.multinomial(probs, num_samples=1)  # (batch_size, 1)

        # Otherwise same as before: get idx of the vocab entry with the highest logits value
        else:
            idx_next = torch.argmax(logits, dim=-1, keepdim=True)  # (batch_size, 1)

        if idx_next == eos_id:  # Stop generating early if end-of-sequence token is encountered and eos_id is specified
            break

        # Same as before: append sampled index to the running sequence
        idx = torch.cat((idx, idx_next), dim=1)  # (batch_size, num_tokens+1)

    return idx


def stream_next_tokens(model, idx, max_new_tokens, context_size, temperature=0.0, top_k=None, eos_id=None):
    """
    Generator version of generate_next_tokens.
    Yields each new token as it is generated.
    """
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -context_size:]
        with torch.no_grad():
            logits = model(idx_cond)
        logits = logits[:, -1, :]

        if top_k is not None:
            top_logits, _ = torch.topk(logits, top_k)
            min_val = top_logits[:, -1]
            logits = torch.where(logits < min_val, torch.tensor(float("-inf")).to(logits.device), logits)

        if temperature > 0.0:
            logits = logits / temperature
            logits = logits - logits.max(dim=-1, keepdim=True).values
            probs = torch.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
        else:
            idx_next = torch.argmax(logits, dim=-1, keepdim=True)

        if idx_next == eos_id:
            break

        idx = torch.cat((idx, idx_next), dim=1)
        yield idx_next